Avoid converting chunked data to Numpy arrays (i.e. .values calls) - #1588
Open
cmdupuis3 wants to merge 20 commits into
Open
Avoid converting chunked data to Numpy arrays (i.e. .values calls)#1588cmdupuis3 wants to merge 20 commits into
.values calls)#1588cmdupuis3 wants to merge 20 commits into
Conversation
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Collaborator
Author
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Collaborator
Author
|
There's some spurious asv results from machine variability, but the benchmarks consistenly show peak-mem reductions and some speedups for cross-sections. |
Collaborator
Author
|
pre-commit.ci autofix |
Sevans711
self-requested a review
July 27, 2026 16:18
Sevans711
requested changes
Jul 27, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Hi @cmdupuis3, thank you for proposing these changes! Overall these look like good clean changes which should help improve the scalability of uxarray.
I have some notes/suggestions/requested changes. Primarily, I noticed that there are missing regression tests. This may be especially important for places with significant changes (more significant than just replacing obj.values with obj.data):
UxDataArray.integrate()UxDataArrayCrossSectionAccessor.__call__()uxarray.plot.matplotlib._nearest_neighbor_resample()RemapAccessor.apply_weights()
It may be nice to add a regression test for UxDataset.to_xarray() too, even though the changes here are minor, since that feels like a very core part of the functionality.
The tests could be similar to what you added already in test_topological_agg.py (plus my additional request on that file): ensure that numpy and dask inputs ultimately give the same values (plus assert that the dask inputs lead to dask outputs).
Other parts of the code changes here might benefit from similar tests, but I don't know if that should be necessary. For example, the changes in _geos just replace obj.values with obj.data; does a regression test need to be added for that or no? Curious to hear from @erogluorhan and/or @rajeeja on this question in particular.
Misc. note: I'm not sure if I fully understand the changes to uxarray/cross_sections files. I need to make sure to take a closer look at those during a subsequent review.
| "i,...i", self.uxgrid.face_areas.values, self.values | ||
| ) | ||
| else: | ||
| # dask-backed data: xr.dot keeps the reduction lazy |
Collaborator
There was a problem hiding this comment.
Does integration need a regression test which is sensitive to dask versus numpy? Perhaps, assert that the result is a dask array, and that the values equal the numpy version, similar to the tests already added by this PR into test_topological_agg.
erogluorhan
changed the title
.values calls) in the code
erogluorhan
changed the title
.values calls) in the code.values calls)
Collaborator
Author
|
pre-commit.ci autofix |
Collaborator
Author
|
@erogluorhan Sam and I were wondering about your philosophy with respect to dask versus numpy routines. I was thinking that if we have a case where the numpy routine could be faster than a dask routine in some situations, we might want to keep both, or maybe you'd want to go all in on dask for simplicity. if isinstance(uxda.data, np.ndarray):
aggregated_var = _apply_node_to_face_aggregation_numpy(uxda, ...)
elif isinstance(uxda.data, da.Array):
# apply aggregation on dask array, TODO:
aggregated_var = _apply_node_to_face_aggregation_numpy(uxda, ...)
else:
raise ValueErrorThis is something in main's aggregate.py. I fleshed out the dask branch, but now we still have this logic, and I left the numpy routine in on this branch. |
Member
If I am understanding this correctly, question is about the use of routines, not data structures. If so, using standard NumPy functions directly on chunked Dask-backed Xarray objects should be completely safe with chunks. Let me know if this helps. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Partly addresses #1583
Overview
This PR is to resolve suboptimal usage of
.valuesthroughout the repo, primarily by deferring to lazy xarray and dask operations. See the table of usage sites in Issue #1583.This specific PR is for the first two tables of usage sites; dead code and bugfixes aren't really included here. Aside from simple cases, some cases were solved by branching on whether the data type was already chunked, so scalability for some parts will depend on if you're using numpy or xarray/dask arrays at those points.
PR Checklist
General
Testing
Documentation